#include "router.h"
#include "tunnel.h"
#include "../common/codec.h"
#include "../common/factory.h"
#include "../common/console.h"
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <net/if.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <linux/if_ether.h>
#include <linux/if_packet.h>



Router::Router (const Config& cfg) : _cfg(cfg) {
  TRACE
  abort_assert(cfg.min_port > 1024 && cfg.max_port <= 65535);
  auto base = _cfg.min_port;
  auto size = _cfg.max_port - _cfg.min_port + 1;
  _rules = new Rule[size];
  for (auto i = 0; i < size; ++i) {
    _rules[i].port = base++;
    _chain.push_back(&_rules[i].joint);
  }
  _focus = ntohl(inet_addr(_cfg.focus));

  _socket = socket(AF_PACKET, SOCK_DGRAM, htons(ETH_P_IP));
  abort_assert(_socket > 0);

  struct ifreq ifr;
  ifr.ifr_addr.sa_family = AF_INET;
  strcpy(ifr.ifr_name, _cfg.if_name);
  abort_assert(-1 != ioctl(_socket, SIOCGIFADDR, &ifr));
  _ipv4 = ntohl(((sockaddr_in *)&ifr.ifr_addr)->sin_addr.s_addr);

  struct sockaddr_ll addr = {0};
  addr.sll_family = AF_PACKET;
  addr.sll_protocol = htons(ETH_P_IP);
  addr.sll_ifindex = if_nametoindex(_cfg.if_name);
  abort_assert(-1 != bind(_socket, (sockaddr*)&addr, sizeof(addr)));
}


Router::~Router () {
  TRACE
}


void Router::startup () {
  TRACE
  Console::run("iptables -t filter -A OUTPUT -p tcp --sport %d:%d --tcp-flags ALL RST -j DROP", _cfg.min_port, _cfg.max_port);
  _upWorker = new std::thread(&Router::upLoop, this);
  _upWorker->detach();
  _downWorker = new std::thread(&Router::downLoop, this);
  _downWorker->detach();
}


void Router::shutdown () {
  TRACE
  delete _upWorker;
  delete _downWorker;
  Console::run("iptables -t filter -D OUTPUT -p tcp --sport %d:%d --tcp-flags ALL RST -j DROP", _cfg.min_port, _cfg.max_port);
  close(_socket);
}


void Router::upLoop () {
  TRACE
  struct sockaddr_ll addr = {0};
  addr.sll_family = AF_PACKET;
  addr.sll_protocol = htons(ETH_P_IP);
  addr.sll_ifindex = if_nametoindex(_cfg.if_name);
  addr.sll_halen = 6;
  memcpy(addr.sll_addr, _cfg.if_addr, 6);
  for (;;) {
    auto pkg = _queue.poll();
    auto &buf = pkg->buf;
    if (pkg->dstIPv4 == _focus)
      log_e("Focus.UP: %zd", buf.readable());
    auto nwrite = sendto(_socket, buf.readIter(), buf.readable(), 0, (const sockaddr*)&addr, sizeof(addr));
    if (nwrite < 0) {
      log_e("Router: write socket error");
      return;
    }
  }
}


void Router::downLoop () {
  TRACE
  u8 buf[IpPackage::MAX_SIZE];
  for (;;) {
    struct sockaddr_ll addr = {0};
    socklen_t len = sizeof(addr);
    auto nread = recvfrom(_socket, buf, IpPackage::MAX_SIZE, 0, (sockaddr*)&addr, &len);
    if (nread <= 0) {
      log_e("Router: read socket error");
      return;
    }
    _cache_mutex.lock();
    _cache.normalize();
    _cache.write(buf, (uwl)nread);
    _cache_mutex.unlock();
    Factory::executor().execute([this] () {
      for (auto i = 0; i < 10; ++i) {
        _cache_mutex.lock();
        auto pkg = Codec::parse(_cache);
        _cache_mutex.unlock();
        if (!pkg)
          break;
        if (pkg->dstIPv4 != _ipv4 || pkg->dstPort < _cfg.min_port || pkg->dstPort > _cfg.max_port)
          continue;
        if (!pkg->err.empty()) {
          log_w("Router.package: %s", pkg->err.c_str());
          continue;
        }
        if (pkg->srcIPv4 == _focus)
          log_i("Focus.DOWN: %zd", pkg->buf.readable());
        forwardDown(pkg);
      }
    });
  }
}


void Router::forwardUp (Ptr<Tunnel> tunnel, Ptr<IpPackage> pkg) {
  TRACE
  _rules_mutex.lock();
  auto rule = getRuleBySrc(std::move(tunnel), pkg->srcIPv4, pkg->srcPort, pkg->dstIPv4, pkg->dstPort);
  pkg->srcIPv4 = _ipv4;
  pkg->srcPort = rule->port;
  _rules_mutex.unlock();
  Codec::nat(pkg);
  _queue.push(std::move(pkg));
}


bool Router::forwardDown (Ptr<IpPackage> pkg) {
  TRACE
  _rules_mutex.lock();
  auto rule = getRuleByDst(pkg->srcIPv4, pkg->srcPort, pkg->dstIPv4, pkg->dstPort);
  if (!rule) {
    _rules_mutex.unlock();
    log_e("No rule for %s -> %s", Codec::iptcpaddr(pkg->srcIPv4, pkg->srcPort).c_str(), Codec::iptcpaddr(pkg->dstIPv4, pkg->dstPort).c_str());
    return false;
  }
  pkg->dstIPv4 = rule->srcIPv4;
  pkg->dstPort = rule->srcPort;
  auto tunnel = rule->tunnel;
  _rules_mutex.unlock();
  Codec::nat(pkg);
  tunnel->write(pkg);
  return true;
}


Rule* Router::getRuleBySrc (Ptr<Tunnel> tunnel, u32 srcIPv4, u16 srcPort, u32 dstIPv4, u16 dstPort) {
  auto key = Rule::srcKey(srcIPv4, srcPort);
  auto iter = _mapper.find(key);
  if (iter != _mapper.end()) {
    auto rule = iter->second;
//    log_v("Set Rule: %s <-> %d <-> %s", Codec::printAddress(srcIPv4, srcPort).c_str(), rule->port, Codec::printAddress(dstIPv4, dstPort).c_str());
    rule->dstIPv4 = dstIPv4;
    rule->dstPort = dstPort;
    rule->tunnel = std::move(tunnel);
    active(rule);
    return rule;
  }
  else {
    auto rule = _chain.next()->cast(&Rule::joint);
    log_v("Add Rule: %s <-> %d <-> %s", Codec::iptcpaddr(srcIPv4, srcPort).c_str(), rule->port, Codec::iptcpaddr(dstIPv4, dstPort).c_str());
    rule->srcIPv4 = srcIPv4;
    rule->srcPort = srcPort;
    rule->dstIPv4 = dstIPv4;
    rule->dstPort = dstPort;
    rule->tunnel = std::move(tunnel);
    _mapper.emplace(key, rule);
    active(rule);
    return rule;
  }
}


Rule* Router::getRuleByDst (u32 srcIPv4, u16 srcPort, u32 dstIPv4, u16 dstPort) {
  if (dstIPv4 != _ipv4) {
    log_e("dstIPv4 not matches");
    return nullptr;
  }
  if (dstPort < _cfg.min_port || _cfg.max_port < dstPort) {
    log_e("dstPort not matches");
    return nullptr;
  }
  auto rule = &_rules[dstPort - _cfg.min_port];
  if (srcIPv4 != rule->dstIPv4) {
    log_e("srcIPv4 not matches");
    return nullptr;
  }
  if (srcPort != rule->dstPort) {
    log_e("srcPort not matches");
    return nullptr;
  }
  if (!rule->tunnel) {
    log_e("rule invalid");
    return nullptr;
  }
  active(rule);
  return rule;
}


void Router::releaseMapping (Ptr<Tunnel> tunnel, u32 srcIPv4, u16 srcPort) {
  TRACE
  std::lock_guard<std::mutex> lock {_rules_mutex};
  auto iter = _mapper.find(Rule::srcKey(srcIPv4, srcPort));
  if (iter != _mapper.end()) {
    auto rule = iter->second;
    if (rule->tunnel == tunnel) {
//    log_i("Del Rule: %s <-> %d <-> %s", Codec::printAddress(srcIPv4, srcPort).c_str(), rule->port, Codec::printAddress(rule->dstIPv4, rule->dstPort).c_str());
      rule->srcIPv4 = 0;
      rule->srcPort = 0;
      rule->dstIPv4 = 0;
      rule->dstPort = 0;
      rule->tunnel = nullptr;
      inactive(rule);
      _mapper.erase(iter);
    }
  }
}


void Router::printRuleChain () {
  for (auto iter = _chain.next(); iter != &_chain; iter = iter->next()) {
    auto rule = iter->cast(&Rule::joint);
    if (rule->tunnel)
      log_d("%s <-> %d <-> %s", Codec::iptcpaddr(rule->srcIPv4, rule->srcPort).c_str(), rule->port, Codec::iptcpaddr(rule->dstIPv4, rule->dstPort).c_str());
  }
}
